# 清理环境，避免重复记录
# del openai

from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

from langchain.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough

model = ChatOpenAI(model="gpt-4o")

prompt = ChatPromptTemplate.from_messages([
    HumanMessagePromptTemplate.from_template("Say hello to {input}!")
])


# 定义输出解析器
parser = StrOutputParser()

chain = (
    {"input": RunnablePassthrough()}
    | prompt
    | model
    | parser
)

from langfuse.decorators import langfuse_context, observe

@observe()
def run():
    langfuse_context.update_current_trace(
            name="LangChainDemo",
            user_id="wzp",
        )
    
    # 获取当前 LangChain 回调处理器
    langfuse_handler = langfuse_context.get_current_langchain_handler()
    
    return chain.invoke(input="AGIClass", config={"callbacks": [langfuse_handler]})

print(run())
langfuse_context.flush() # Langfuse 回传记录是异步的，可以通过 flush 强制更新

